import torch
from math import sqrt


FEATMAP_DIMS = (38, 19, 10, 5, 3, 1)

OBJECT_SCALES = (0.1, 0.2, 0.375, 0.55, 0.725, 0.9)

ASPECT_RATIOS = ((1.0, 2.0, 0.5), 
                 (1.0, 2.0, 3.0, 0.5, 0.333), 
                 (1.0, 2.0, 3.0, 0.5, 0.333),
                 (1.0, 2.0, 3.0, 0.5, 0.333),
                 (1.0, 2.0, 0.5), 
                 (1.0, 2.0, 0.5))

device = 'cuda' if torch.cuda.is_available() else 'cpu'


class SSDAnchorCreator(object):
    """
    create anchor boxes over all scales of feature maps
        each anchor is in cxcywh format.
    """
    def __init__(self, featmap_dims=FEATMAP_DIMS, 
                 object_scales=OBJECT_SCALES, 
                 aspect_ratios=ASPECT_RATIOS):
        self.featmap_dims = featmap_dims
        self.object_scales = object_scales
        self.aspect_ratios = aspect_ratios
    
    def create(self):
        anchors = []
        for k, dim in enumerate(self.featmap_dims):
            scale = self.object_scales[k]

            for i in range(dim):
                for j in range(dim):
                    cx = (j + 0.5) / dim
                    cy = (i + 0.5) / dim
                    for ratio in self.aspect_ratios[k]:
                        anchor = [cx, cy, scale * sqrt(ratio), scale / sqrt(ratio)]
                        anchors.append(anchor)

                        if ratio == 1:
                            try:
                                additional_scale = sqrt(scale * self.object_scales[k + 1])
                            except IndexError:
                                additional_scale = 1.0
                            anchor = [cx, cy, additional_scale, additional_scale]
                            anchors.append(anchor)
        anchors = torch.FloatTensor(anchors).to(device)
        anchors.clamp_(0, 1)
        return anchors
